from langchain.chains import GraphCypherQAChain
from langchain_community.graphs import Neo4jGraph
from langchain_community.chat_models import QianfanChatEndpoint
from langchain.prompts.prompt import PromptTemplate
import os

# 设置百度千帆 API 密钥
os.environ["QIANFAN_ACCESS_KEY"] = "ALTAKiwnHEggAvYXbrd21K3luc"
os.environ["QIANFAN_SECRET_KEY"] = "0d9cd08f6a924dac8673281d172ba485"

graph = Neo4jGraph(
    url="bolt://localhost:7687",
    username="neo4j",
    password="12345678"
)

llm = QianfanChatEndpoint(model="ERNIE-Bot")

cypher_prompt = PromptTemplate.from_template("""
你是一个擅长图数据库查询的专家，请根据问题用 Cypher 查询语言生成查询语句，要求：
1. 使用 MATCH 和 OPTIONAL MATCH，而不是 size() 函数。
2. 所有关系使用明确变量，例如 [r]，并用 COUNT 来统计数量。
3. 不允许在 size() 内直接用模式表达式。
4. 返回结果字段应包含实体名称和关系强度（如参与者数量）。
5. 查询目标：{question}

请仅返回合法 Cypher 语句，不要加自然语言解释。
""")

chain = GraphCypherQAChain.from_llm(
    llm,
    graph=graph,
    verbose=True,
    allow_dangerous_requests=True
)

# 执行查询（使用 invoke 替代 run）
try:
    response = chain.invoke({"query": "董卓和吕布的关系"})
    print(response)
    print(type(response))
except Exception as e:
    print(f"查询失败: {str(e)}")

""" response = chain.run("赤壁之战参战方实力分析") """


